from core.function import Function
from core.cuda import cuda_module
from core.config import Config
import utils

# Batch Normalization（批标准化）是一种用于深度神经网络中的优化技术。
# 其主要目的是加速神经网络的训练过程，减少梯度消失问题，加速收敛，同时提高模型的泛化能力。

# Batch Normalization 的操作如下：
#
# 1. 对每个 mini-batch 的输入数据进行标准化，即减去均值并除以标准差。
# 2. 对标准化后的数据进行缩放和平移，通过缩放参数（gamma）和平移参数（beta）进行调整。
# 3. 这样可以确保每一层的输入分布在训练过程中始终保持稳定，有助于加速网络的收敛。

class BatchNorm(Function):
    def __init__(self, mean, var, decay, eps):
        """
        初始化 BatchNorm 函数。

        Parameters:
            mean (Variable): 平均值。
            var (Variable): 方差。
            decay (float): 移动平均的衰减率。
            eps (float): 避免除以零的小值。
        """
        self.avg_mean = mean
        self.avg_var = var
        self.decay = decay
        self.eps = eps
        self.inv_std = None

    def forward(self, x, gamma, beta):
        """
        执行 BatchNorm 操作的前向传播。

        Parameters:
            x (Variable): 输入变量。
            gamma (Variable): 缩放参数。
            beta (Variable): 偏移参数。

        Returns:
            Variable: 操作结果。
        """
        assert x.ndim == 2 or x.ndim == 4

        x_ndim = x.ndim
        if x_ndim == 4:
            N, C, H, W = x.shape
            # (N, C, H, W) -> (N*H*W, C)
            x = x.transpose(0, 2, 3, 1).reshape(-1, C)

        xp = cuda_module

        if Config.train:
            mean = x.mean(axis=0)
            var = x.var(axis=0)
            inv_std = 1 / xp.sqrt(var + self.eps)
            xc = (x - mean) * inv_std

            m = x.size // gamma.size
            s = m - 1. if m - 1. > 1. else 1.
            adjust = m / s  # unbiased estimation
            self.avg_mean *= self.decay
            self.avg_mean += (1 - self.decay) * mean
            self.avg_var *= self.decay
            self.avg_var += (1 - self.decay) * adjust * var
            self.inv_std = inv_std
        else:
            inv_std = 1 / xp.sqrt(self.avg_var + self.eps)
            xc = (x - self.avg_mean) * inv_std
        y = gamma * xc + beta

        if x_ndim == 4:
            # (N*H*W, C) -> (N, C, H, W)
            y = y.reshape(N, H, W, C).transpose(0, 3, 1, 2)
        return y

    def backward(self, gy):
        """
        执行 BatchNorm 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            tuple: 相对于输入变量的梯度，缩放参数的梯度，偏移参数的梯度。
        """

        gy_ndim = gy.ndim
        if gy_ndim == 4:
            N, C, H, W = gy.shape
            gy = gy.transpose(0, 2, 3, 1).reshape(-1, C)

        x, gamma, beta = self.inputs
        batch_size = len(gy)

        if x.ndim == 4:
            N, C, H, W = x.shape
            x = x.transpose(0, 2, 3, 1).reshape(-1, C)
        mean = x.sum(axis=0) / batch_size
        xc = (x - mean) * self.inv_std

        gbeta = utils.functions_collect.sum(gy, axis=0)
        ggamma = utils.functions_collect.sum(xc * gy, axis=0)
        gx = gy - gbeta / batch_size - xc * ggamma / batch_size
        gx *= gamma * self.inv_std

        if gy_ndim == 4:
            gx = gx.reshape(N, H, W, C).transpose(0, 3, 1, 2)
        return gx, ggamma, gbeta

    def batch_norm(self, x, gamma, beta):
        """
        执行 BatchNorm 操作。

        Parameters:
            x (Variable): 输入变量。
            gamma (Variable): 缩放参数。
            beta (Variable): 偏移参数。

        Returns:
            Variable: 操作结果。
        """
        return self(x, gamma, beta)

